package com.shen.langchain4j.controller;

import com.shen.langchain4j.service.ChatAssistant;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * 大模型控制层接口
 */
@Slf4j
@RestController
@RequestMapping(value = "/streamChat")
public class StreamingChatController {
    /**
     * 低阶的LLM API
     */
    @Autowired
    @Qualifier("streamQwen")
    private StreamingChatModel streamQwen;
    /**
     * 高阶的LLM API(自己封装接口)
     */
    @Autowired
    @Qualifier("chatAssistant")
    private ChatAssistant chatAssistant;

    /**
     * 低阶LLM API 流式返回到页面
     * @param prompt 提示词
     * @return 异步序列
     */
    @GetMapping(value = "/chatOne")
    public Flux<String> chatOne(@RequestParam(value = "prompt", defaultValue = "你是谁") String prompt) {
        return Flux.create(emitter -> streamQwen.chat(prompt, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                emitter.next(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                emitter.complete();
            }

            @Override
            public void onError(Throwable error) {
                emitter.error(error);
            }
        }));
    }

    /**
     * 低阶LLM API 流式返回到后端
     * @param prompt 提示词
     */
    @GetMapping(value = "/chatTwo")
    public void chatTwo(@RequestParam(value = "prompt", defaultValue = "你是谁") String prompt) {
        streamQwen.chat(prompt, new StreamingChatResponseHandler() {

            @Override
            public void onPartialResponse(String partialResponse) {
                System.out.println(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                System.out.println("==response over:" + completeResponse);
            }

            @Override
            public void onError(Throwable error) {
                error.printStackTrace();
            }
        });
    }

    /**
     * 高阶LLM API 直接调用封装接口流式返回到页面
     * @param prompt 提示词
     * @return 异步序列
     */
    @GetMapping(value = "/chatThree")
    public Flux<String> chatThree(@RequestParam(value = "prompt", defaultValue = "你是谁") String prompt) {
        return chatAssistant.chatFlux(prompt);
    }

}
